import os
from typing import Sequence, Iterable

import numpy as np
import pandas as pd
from sklearn.metrics import classification_report

class Report:
    def __init__(self, types:Sequence=None, indicators:Sequence=None) -> None:
        '''
        Args:
            types: Extract the types of scores.
            indicators: Selected indicators.
        '''
        self._types = types or ('acc', 'macro', 'weighted')  # accuracy/macro avg/weighted avg
        self._indicators = indicators or ('precision', 'recall', 'f1-score')
    
    def set_types(self, *types):
        self._types = types
        return self
    
    def set_indicators(self, *indicators):
        self._indicators = indicators
        return self
    
    def extract_scores(self, y_true, y_pred, one_hot:bool=False):
        '''提取分数
        :param y_true: 真实值
        :param y_pred: 预测值
        :param one_hot: 是否为独热编码
        '''
        self.scores = []
        if one_hot:
            y_true = np.argmax(y_true, axis=-1)
            y_pred = np.argmax(y_pred, axis=-1)
        report = classification_report(y_true, y_pred, zero_division=0, output_dict=True)
        for _type in self._types:

            if _type in 'accuracy':
                self.scores.append(report['accuracy'])
                continue

            for k in report.keys():
                if _type in k:
                    self.scores.append(list(report[k].values())[:len(self._indicators)])
                    break
            else:
                self.scores.append(['nan']*len(self._indicators))
        return self
    
    def make_report(self, sep:str=' '):
        '''制作报告
        :param sep: 类型和指标的连接符
        '''
        self.report = pd.DataFrame(columns=['indicator', 'score'])
        for _type, _scores in zip(self._types, self.scores):

            if _type in 'accuracy':
                self.report.loc[len(self.report)] = [_type, _scores]
                continue

            indicators = [f'{_type}{sep}{indicator}' for indicator in self._indicators]
            for indicator, score in zip(indicators, _scores):
                self.report.loc[len(self.report)] = [indicator, score]
        return self

    def extend_report(self, field_name, field_value, insert:int=None):
        '''扩展报告
        :param field_name: 字段名
        :param field_value: 字段值
        :param insert: 插入报告列索引位置
        '''
        insert = self.report.shape[1] if insert is None else insert
        if field_name in self.report.columns:
            insert = self.report.columns.to_list().index(field_name)
            self.report.drop(columns=[field_name], inplace=True)
        self.report.insert(insert, field_name, field_value)
        return self
    
    def show_report(self, fields=None):
        '''展示报告
        :param fields: 展示特定字段
        '''
        if fields is None:
            print(self.report)
        else:
            print(self.report[fields])
        return self
    
    def save_report(self, path, subset=None):
        '''保存报告
        :param subset: 删除指定子集(以列为单位)中含重复值的行
        '''
        if os.path.exists(path):
            df = pd.read_csv(path)
            self.report = pd.concat([df, self.report])
            self.report.drop_duplicates(inplace=True, subset=subset, keep='last')
        self.report.to_csv(path, index=None)

def auto_report(y_true=None, y_pred=None, scores:Sequence=None, 
                types:Sequence=None, indicators:Sequence=None, 
                extend_dict:dict=None, save_path:str=None, 
                verbose:bool=True) -> None:
    '''自动报告
    :param y_true: 真实值
    :param y_pred: 预测值
    :param scores: 自定义分数
    :param extend_dict: 扩展报告字典
    :param save_path: 报告保存路径
    :param verbose: 是否展示详细信息
    '''
    report = Report()
    if types is not None:
        report.set_types(*types)
    if indicators is not None:
        report.set_indicators(*indicators)
    if y_pred is not None and y_true is not None:
        report.extract_scores(y_true, y_pred)
    if scores is not None:
        if isinstance(scores, Iterable):
            scores = [scores]
        report.scores = scores
    report.make_report()
    if extend_dict is not None:
        for key, value in extend_dict.items():
            report.extend_report(key, value)
    if verbose:
        report.show_report()
    if save_path is not None:
        columns:pd.Index = report.report.columns
        columns = columns.drop('score')
        report.save_report(path=save_path, subset=columns)
